package com.cpp.server.init;

import jakarta.annotation.PostConstruct;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.document.Document;
import org.springframework.ai.reader.ExtractedTextFormatter;
import org.springframework.ai.reader.pdf.PagePdfDocumentReader;
import org.springframework.ai.reader.pdf.config.PdfDocumentReaderConfig;
import org.springframework.ai.transformer.splitter.TextSplitter;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.io.Resource;
import org.springframework.core.io.support.ResourcePatternResolver;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.stereotype.Component;

import java.io.File;
import java.io.InputStream;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Set;

// Apache POI imports
import org.apache.poi.hwpf.HWPFDocument;
import org.apache.poi.xwpf.usermodel.XWPFDocument;
import org.apache.poi.xwpf.usermodel.XWPFParagraph;

@Component
@Slf4j
@RequiredArgsConstructor
public class Init {

    private final VectorStore vectorStore;
    private final ResourcePatternResolver resourcePatternResolver;
    private final StringRedisTemplate stringRedisTemplate;

    // 使用 TokenTextSplitter 对文档内容进行切分
    private final TextSplitter textSplitter = new TokenTextSplitter(); // 每个段落最多 512 tokens
    @Value("${config.file.path}")
    private  String path;

//    @PostConstruct
    public void init() {
//        redis查询doc*开头的key
        Set<String> docKeys = stringRedisTemplate.keys("doc*");
        if (docKeys != null && !docKeys.isEmpty()) {
            log.info("🔍 发现以下 doc* 开头的 key：{}", docKeys);
            //redis有向量数据不重新读取
            return;
        }
        try {
            // 加载 classpath:classResources/ 下的所有文件
            Resource[] resources = resourcePatternResolver.getResources(path);

            for (Resource resource : resources) {
                String filename = resource.getFilename();
                if (filename == null) continue;

                log.info("📄 正在处理文件：{}", filename);

                List<Document> documents = new ArrayList<>();

                if (filename.endsWith(".pdf")) {
                    documents.addAll(readPdf(resource));
                } else if (filename.endsWith(".doc")) {
                    documents.addAll(readDoc(resource));
                } else if (filename.endsWith(".docx")) {
                    documents.addAll(readDocx(resource));
                } else {
                    log.warn("🚫 不支持的文件格式：{}", filename);
                    continue;
                }
                for (Document document : documents) {
                    vectorStore.add(Collections.singletonList(document));
                }
                log.info("✅ {} 共提取 {} 条 Document 并已写入向量库", filename, documents.size());
            }

            log.info("🎉 初始化完成，所有文档已导入向量库");

        } catch (Exception e) {
            log.error("❌ 初始化失败", e);
        }
    }

    // ===================== 文件读取方法 =====================

    private List<Document> readPdf(Resource resource) {
        PagePdfDocumentReader reader = new PagePdfDocumentReader(
                resource,
                PdfDocumentReaderConfig.builder()
                        .withPageExtractedTextFormatter(ExtractedTextFormatter.defaults())
                        .withPagesPerDocument(1)
                        .build()
        );
        List<Document> documents = reader.read();

        // 对每页内容进一步切分
        List<Document> splitDocuments = new ArrayList<>();
        for (Document doc : documents) {
            splitDocuments.addAll(textSplitter.split(doc));
        }

        return splitDocuments;
    }

    private List<Document> readDoc(InputStream is) throws Exception {
        HWPFDocument wdoc = new HWPFDocument(is);
        String fullText = wdoc.getDocumentText();
        Document  doc = new Document(fullText);

        // 按换行符分割为段落
        return textSplitter.split(doc);
    }

    private List<Document> readDoc(Resource resource) throws Exception {
        try (InputStream is = resource.getInputStream()) {
            return readDoc(is);
        }
    }

    private List<Document> readDocx(InputStream is) throws Exception {
        XWPFDocument doc = new XWPFDocument(is);
        List<Document> documents = new ArrayList<>();

        for (XWPFParagraph paragraph : doc.getParagraphs()) {
            String text = paragraph.getText();
            if (text != null && !text.trim().isEmpty()) {
                documents.add(new Document(text));
            }
        }

        return documents;
    }

    private List<Document> readDocx(Resource resource) throws Exception {
        try (InputStream is = resource.getInputStream()) {
            return readDocx(is);
        }
    }
}
